#!/usr/bin/python
#
# iphoto_export.py - export albums from iphoto
#
# Copyright (C) 2011-2012 by Johannes Overmann <Johannes.Overmann@gmx.de>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.



import optparse
import xml.etree.ElementTree
import os
import sys
import fnmatch
import shutil
import exceptions
import unicodedata
import hashlib
import codecs

options = None # set in main

# --- helpers --------------------------------------------------------------------------------

class Error(exceptions.Exception):
    """Simple error exception.
    
    Carries a message and a path.
    """
    def __init__(self, message, path = None):
        self.message = message
        self.path = path

        
def latin1(anyStr):
    """Convert string (str or unicode) to latin1 bytestring.
    """
    return unicodedata.normalize("NFKC", unicode(anyStr)).encode("latin-1")


def makeDirs(path):
    """Create dirs recursively.

    If the path already exists this function does nothing.
    """
    if os.path.isdir(path):
        return
    elif not os.path.exists(path):
        os.makedirs(path)
    else:
        raise Error("non-dir already exists while trying to create dir '%s'" % path)

    
# --- plist XML parser --------------------------------------------------------------------------------

def checkTag(tree, expectedTag, path):
    """Check whether node tree is tagged with expectedTag.
    
    Else raise an Error exception.
    """
    if tree.tag != expectedTag:
        raise Error("error: expected tag '%s'" % expectedTag, path)
    

def decodePlist(tree, path):
    """Decode Apple property list (top level).
    
    Return a dict().
    """
    checkTag(tree, "plist", path)
    path += ".plist"
    if (len(list(tree)) != 1) or (list(tree)[0].tag != "dict"):
        raise Error("expecting exactly one dict below 'plist'", path)
    return decodeDict(list(tree)[0], path)


def decodeDict(tree, path):
    """Decode dict from Apple property list.
    
    Return a python dict().
    """
    checkTag(tree, "dict", path)
    path += ".dict"
    r = {}
    index = 0
    for i in list(tree):
        if (index & 1) == 0:
            key = decodeKey(i, path)
        else:
            if i.tag == "key":
                raise Error("not expecting key after key in dict", path)
            value = decodeAny(i, path + "." + key)
            r[key] = value
        index += 1
    return r


def decodeArray(tree, path):
    """Decode array from Apple property list.
    
    Return a python list.
    """
    checkTag(tree, "array", path)
    path += ".array"
    r = []
    index = 0
    for i in list(tree):
        value = decodeAny(i, path + ("[%d]" % index))
        r.append(value)
        index += 1
    return r


def decodeKey(tree, path):
    """Decode key from Apple property list.
    
    Return a string.
    """
    checkTag(tree, "key", path)
    return tree.text


def decodeString(tree, path):
    """Decode string from Apple property list.
    
    Return a string.
    """
    checkTag(tree, "string", path)
    return tree.text


def decodeInteger(tree, path):
    """Decode integer from Apple property list.
    
    Return an int.
    """
    checkTag(tree, "integer", path)
    return int(tree.text)

def decodeReal(tree, path):
    """Decode real from Apple property list.
    
    Return a float.
    """
    checkTag(tree, "real", path)
    return float(tree.text)


def decodeAny(tree, path):
    """Decode any node type depending on the tag.
    
    Returns the appropriate python type.
    """
    if tree.tag == "plist":
        return decodePlist(tree, path)
    elif tree.tag == "dict":
        return decodeDict(tree, path)
    elif tree.tag == "array":
        return decodeArray(tree, path)
    elif tree.tag == "key":
        return decodeKey(tree, path)
    elif tree.tag == "string":
        return decodeString(tree, path)
    elif tree.tag == "integer":
        return decodeInteger(tree, path)
    elif tree.tag == "real":
        return decodeReal(tree, path)
    elif tree.tag == "true":
        return True
    elif tree.tag == "false":
        return False
    else:
        print "error: decodeNode(): unknown tag '%s' (%s)" % (tree.tag, path)
#        raise Error("decodeNode(): unknown tag '%s'" % tree.tag, path)


def dumpTree(tree, depth, indent = 0):
    """Dump Apple property list tree in tag/text format.
    """
    print "%s%s: %s" % ("  " * indent, tree.tag, repr(tree.text))
    if depth > 0:
        for i in list(tree):
            dumpTree(i, depth - 1, indent + 1)


# --- main functionality --------------------------------------------------------------------------------    

def prependFolderName(album, albumIdToAlbum):
    """Recursively prepend the parent folder names to AlbumName.
    Modifies album["AlbumName"].
    """
    if "__has_expanded_folders__" in album:
        return
    album["__has_expanded_folders__"] = True # avoid endless recursion on loops
    if "Parent" not in album:
        return
    parentId = album["Parent"]
    if parentId not in albumIdToAlbum:
        raise Error("album '%s': parent album with id '%s' not found" % (latin1(album["AlbumName"]), parentId))
    parentAlbum = albumIdToAlbum[parentId]
    # recusively prepare album name of parent
    prependFolderName(parentAlbum, albumIdToAlbum)
    # prepend parent Album name as folder name
    album["AlbumName"] = parentAlbum["AlbumName"] + "/" + album["AlbumName"]
        

def addHierarchyToAlbumName(data):
    """Prepend folder names to AlbumName of all albums.
    """
    # pass 1: get dict from albumId to album
    albumList = data["List of Albums"]
    albumIdToAlbum = {}
    for album in albumList:
        if "AlbumId" in album:
            albumIdToAlbum[album["AlbumId"]] = album
    
    # pass 2: recursively prepend folder names to the AlbumName
    for album in albumList:
        prependFolderName(album, albumIdToAlbum)

        
class TargetInfo:
    """Information about a single target file.
    """
    def __init__(self, sourceFileNameU, foto):
        """
        """
        self.sourceFileNameU = sourceFileNameU
        self.foto = foto
        self.checksum = None

        
    def getChecksum(self):
        """Return checksum of files content.
        """
        if self.checksum:
            return self.checksum
        h = hashlib.sha1()
        f = open(self.sourceFileNameU, "rb")
        content = f.read()
        f.close()
        h.update(content)
        self.checksum = h.hexdigest()
        return self.checksum
        

class ExportState:
    """Export state.
    """
    pass

        
def addTarget(exportState, targetFile2targetInfo, lowerTargetFile2targetFile, targetFileName, sourceFileNameU, foto):
    """Add target file to targetFile2targetInfo dict.
    Resolves all target file name conflicts in some way.
    """
    if not os.path.exists(sourceFileNameU):
        print "error: file '%s' not found" % latin1(sourceFileNameU)
        exportState.numErrorsSrcNotFound += 1
        return
    targetInfo = TargetInfo(sourceFileNameU, foto)
    firstCaseTargetFileName = lowerTargetFile2targetFile[targetFileName.lower()]
    if firstCaseTargetFileName not in targetFile2targetInfo:
        # target file does not yet exist: good
        targetFile2targetInfo[firstCaseTargetFileName] = targetInfo
        return
    
    # target file already exists
    # check whether the file has the same content
    firstTargetInfo = targetFile2targetInfo[firstCaseTargetFileName]
    albumName = latin1(os.path.dirname(firstCaseTargetFileName))
    if firstTargetInfo.getChecksum() == targetInfo.getChecksum():
        if options.verbose:
            print "warning: album '%s': files '%s' and '%s' have the same content: ignoring the latter one" % (albumName, latin1(firstTargetInfo.sourceFileNameU), latin1(targetInfo.sourceFileNameU))
        exportState.numClashSameContent += 1
        # we do not add this target file to targetFile2targetInfo since it already exists with the same content
        return 

    # generate new unique target filename by appending _2 etc before the extension
    for i in xrange(2, 1000000):
        nameExt = os.path.splitext(targetFileName)
        newTargetFileName = nameExt[0] + ("_%d" % i) + nameExt[1]
        if newTargetFileName.lower() not in lowerTargetFile2targetFile:
            break
    targetFile2targetInfo[newTargetFileName] = targetInfo
    exportState.numClashDifferentContent += 1
    if options.verbose:
        print "warning: album '%s': renaming file '%s' to '%s' since the former already is used in this album (from '%s') and they differ in content" % (albumName, latin1(targetInfo.sourceFileNameU), latin1(os.path.basename(newTargetFileName)), latin1(firstTargetInfo.sourceFileNameU))


def exportAlbums(data):
    """Export iPhoto albums.
    """
    # data.keys() = ['Minor Version', 'List of Keywords', 'Major Version', 'Master Image List', 'Archive Path', 'Application Version', 'ArchiveId', 'List of Rolls', 'List of Faces', 'List of Albums']

    # add hierarchy to album names
    if not options.flatten:
        addHierarchyToAlbumName(data)
    
    # get dict of fotos
    keyToFoto = data["Master Image List"]
    numOriginalLibraryFotos = 0
    for key, foto in keyToFoto.iteritems():
        foto["__used__"] = 0
        if "OriginalPath" in foto:
            numOriginalLibraryFotos += 1
    
    # get number of fotos in all albums
    albumList = data["List of Albums"]
    numMainFotos = 0
    numOriginalFotos = 0
    newAlbumList = []
    for album in albumList:
        albumNameU = album["AlbumName"]
        albumName = latin1(albumNameU)
        albumType = album.get("Album Type", "All Fotos")
        photoCount = album["PhotoCount"]
        fotoKeyList = album["KeyList"]
        if len(fotoKeyList) != photoCount:
            raise Error("album '%s': KeyList contains %d fotos while PhotoCount is %d" % (albumName, len(fotoKeyList), photoCount))
        skipAlbum = (albumType != "Regular") or (photoCount == 0)
        match = False
        for pattern in options.album:
            if fnmatch.fnmatch(albumName, pattern):
                match = True
                break
        if not match:
            skipAlbum = True
        numMainFotosAlbum = 0
        numOriginalFotosAlbum = 0
        for fotoKey in fotoKeyList:
            if fotoKey not in keyToFoto:
                raise Error("album '%s': foto not found for key %s" % (albumName, fotoKey))
            foto = keyToFoto[fotoKey]
            if albumType == "Regular":
                foto["__used__"] += 1
            if not "ImagePath" in foto:
                raise Error("album '%s': ImagePath missing for foto with key %s" % (albumName, fotoKey))
            numMainFotosAlbum += 1
            if "OriginalPath" in foto:
                numOriginalFotosAlbum += 1
            # check for straneg chars in filename which are not in latin-1
            if True:
                try:
                    s = latin1(foto["ImagePath"])
                except exceptions.Exception, e:
                    print "warning: strange chars in filename", repr(foto["ImagePath"]), e
                if "OriginalPath" in foto:
                    try: 
                        s = latin1(foto["OriginalPath"])
                    except exceptions.Exception, e: 
                        print "warning: strange chars in filename", repr(foto["OriginalPath"]), e
                
        if skipAlbum:
            if options.verbose or options.list_albums:
                print "(skiping album type=%-14.14s %5d/%5d %s)" % (albumType, numMainFotosAlbum, numOriginalFotosAlbum, albumName)
            continue
        numMainFotos += numMainFotosAlbum
        numOriginalFotos += numOriginalFotosAlbum
        album["__numOriginalFotos__"] = numOriginalFotosAlbum
        newAlbumList.append(album)
    numTotalFotos = numMainFotos + numOriginalFotos
    # count unused fotos
    numUnusedFotos = 0
    unusedAlbum = {}
    unusedAlbum["KeyList"] = []
    unusedAlbum["AlbumName"] = "Unused Fotos"
    unusedAlbum["Album Type"] = "Unused Fotos"
    for key, foto in keyToFoto.iteritems():
        if foto["__used__"] == 0:
            numUnusedFotos += 1
            if options.list_unused:
                print "unused foto %s" % latin1(foto["ImagePath"])
            if "OriginalPath" in foto:
                numUnusedFotos += 1
                if options.list_unused:
                    print "unused foto %s" % latin1(foto["OriginalPath"])
            unusedAlbum["KeyList"].append(key)
    unusedAlbum["__numOriginalFotos__"] = numUnusedFotos - len(unusedAlbum["KeyList"])
                
    if options.store_unused and numUnusedFotos:
        if options.verbose:
            print "adding album '%s' with %d fotos total (%d/%d)" % (unusedAlbum["AlbumName"], numUnusedFotos, len(unusedAlbum["KeyList"]), numUnusedFotos - len(unusedAlbum["KeyList"]))
        numTotalFotos += numUnusedFotos
        numMainFotos += len(unusedAlbum["KeyList"])
        numOriginalFotos += numUnusedFotos - len(unusedAlbum["KeyList"])
        newAlbumList.append(unusedAlbum)

    lowerTargetFile2targetFile = {}
    for album in newAlbumList:
        if options.verbose or options.list_albums:
            print "matching album type=%-14.14s %5d/%5d %s" % (album["Album Type"], len(album["KeyList"]), album["__numOriginalFotos__"], latin1(album["AlbumName"]))
        albumNameU = album["AlbumName"]
        fotoKeyList = album["KeyList"]
        for fotoKey in fotoKeyList:
            foto = keyToFoto[fotoKey]
            targetFile = albumNameU + "/" + os.path.basename(foto["ImagePath"])
            if targetFile.lower() not in lowerTargetFile2targetFile:
                lowerTargetFile2targetFile[targetFile.lower()] = targetFile
            if "OriginalPath" in foto:
                targetFile = albumNameU + "/" + options.originals_dir + "/" + os.path.basename(foto["OriginalPath"])
                if targetFile.lower() not in lowerTargetFile2targetFile:
                    lowerTargetFile2targetFile[targetFile.lower()] = targetFile
            
    if options.verbose:
        print "--"
        print "%5d matching albums found (of %d)" % (len(newAlbumList), len(albumList))
        print "%5d+%5d=%5d fotos in all matching albums (main+original=total)" % (numMainFotos, numOriginalFotos, numTotalFotos)
        print "%5d+%5d=%5d fotos in foto library (main+original=total)" % (len(keyToFoto), numOriginalLibraryFotos, len(keyToFoto) + numOriginalLibraryFotos)
        print "%5d unused fotos which are not in any regular album" % numUnusedFotos

    # check source and target files consistency
    numProcessedFotos = 0
    targetFile2targetInfo = {}
    exportState = ExportState()
    exportState.numErrorsSrcNotFound = 0
    exportState.numClashSameContent = 0
    exportState.numClashDifferentContent = 0
    for album in newAlbumList:
        albumNameU = album["AlbumName"]
        fotoKeyList = album["KeyList"]
        for fotoKey in fotoKeyList:
            if fotoKey not in keyToFoto:
                raise Error("album '%s': foto missing for key %s" % (latin1(albumNameU), fotoKey))
            foto = keyToFoto[fotoKey]
            addTarget(exportState, targetFile2targetInfo, lowerTargetFile2targetFile, albumNameU + "/" + os.path.basename(foto["ImagePath"]), foto["ImagePath"], foto)
            numProcessedFotos += 1
            if "OriginalPath" in foto:
                addTarget(exportState, targetFile2targetInfo, lowerTargetFile2targetFile, albumNameU + "/" + options.originals_dir + "/" + os.path.basename(foto["OriginalPath"]), foto["OriginalPath"], foto)
                numProcessedFotos += 1                
            if options.verbose:
                sys.stdout.write("checking files: %5d/%5d (%5.1f%%) processed\r" % (numProcessedFotos, numTotalFotos, float(numProcessedFotos)/numTotalFotos*100.0))
                sys.stdout.flush()
    if options.verbose:
        sys.stdout.write("\n")
    if exportState.numErrorsSrcNotFound:
        raise Error("%5d files not found: please correct these path names in file '%s' manually" % (exportState.numErrorsSrcNotFound, options.iphoto_album_data_xml))
    if exportState.numClashSameContent:
        print "%5d files ignored which have the same content as an already existing file in the same album" % (exportState.numClashSameContent)
    if exportState.numClashDifferentContent:
        print "%5d files renamed which differ in content from a file with the same name in the same album" % (exportState.numClashDifferentContent)

    if True:
        f = codecs.open("filelist.txt", encoding = "utf-8", mode = "wb")
        for targetFileName, targetInfo in sorted(targetFile2targetInfo.iteritems()):
            f.write("0_Fotos/" + unicodedata.normalize("NFKD", unicode(targetFileName)));
            f.write("\n");
        f.close()
        
    if options.list_albums:
        return

    # copy files 
    numProcessedFiles = 0
    if options.verbose:
        print "copying %d files ..." % len(targetFile2targetInfo)
    for targetFileName, targetInfo in sorted(targetFile2targetInfo.iteritems()):
        # check whether target file already exists (should not happen)
        fullTargetFilename = options.output_dir + "/" + targetFileName        
        if os.path.exists(fullTargetFilename):
            if not options.update_output:
                print "error: target file '%s' already exists" % (latin1(fullTargetFilename))
        else:
            makeDirs(os.path.dirname(fullTargetFilename))
            shutil.copy2(targetInfo.sourceFileNameU, fullTargetFilename)
        numProcessedFiles += 1
        if options.verbose:
            sys.stdout.write("copying files: %5d/%5d (%5.1f%%) processed\r" % (numProcessedFiles, len(targetFile2targetInfo), float(numProcessedFiles)/len(targetFile2targetInfo)*100.0))
            sys.stdout.flush()
    if options.verbose:
        sys.stdout.write("\n")
        print "done"


def main():
    """Main function of this module.
    """
    global options
    usage = "Usage: %prog [options]"
    version = "0.9.5"
    parser = optparse.OptionParser(usage=usage, version=version)
    parser.add_option("-v", "--verbose", help="verbose output", action="count", default=1)
    parser.add_option("-q", "--quiet",   help="suppress verbose output", action="store_true", default=False)
    parser.add_option("-i", "--iphoto-album-data-xml", type="string", help="iPhoto AlbumData.xml input file (default=%default)", default="~/Pictures/iPhoto Library/AlbumData.xml")
    parser.add_option("-o", "--output-dir", type="string", help="output directory (default=%default)", default="exported_albums")
    parser.add_option("-O", "--originals-dir", type="string", help="name of the subdirectory where the original fotos are stored (choose 'Originals' to have them as invisible originals in Picasa, choose '' to disable exporting them, default=%default)", default="iPhotoOriginals")
    parser.add_option("-a", "--album", action="append", help="album to export (default is to export all albums, fnmatch patterns allowed)", default=[])
    parser.add_option("-l", "--list-albums", action="store_true", help="list available albums and exit", default=False)
    parser.add_option("-U", "--list-unused", action="store_true", help="list unused fotos", default=False)
    parser.add_option("-u", "--store-unused", action="store_true", help="store fotos which are not in any other regular album in album 'Unused Fotos'", default=False)
    parser.add_option("-F", "--flatten", action="store_true", help="ignore folders and export all albums to a flat list of album directories", default=False)
    parser.add_option("-d", "--update-output", action="store_true", help="silently ignore files which already exist in output when copying files, just add new files (no files are removed from output)", default=False)
    (options, args) = parser.parse_args()
    
    # check and modify args
    if options.quiet:
        options.verbose = 0
    if not options.album:
        options.album.append("*")
    options.iphoto_album_data_xml = os.path.expanduser(options.iphoto_album_data_xml)
    if len(args):
        parser.error("no non-option arguments expected")

    # parse AlbumData.xml
    try:
        if options.verbose:
            print "reading iPhoto album data from '%s' ..." % options.iphoto_album_data_xml
        tree = xml.etree.ElementTree.parse(options.iphoto_album_data_xml)
    except exceptions.Exception, e:
        parser.error("error while parsing file '%s': %s" % (options.iphoto_album_data_xml, e))
    #    xml.etree.ElementTree.dump(tree)

    try:
        data = decodeAny(tree.getroot(), "root")
        exportAlbums(data)
    except Error, e:
        if e.path:
            print "error: %s (%s)" % (e.message, e.path)
        else:
            print "error: %s" % (e.message)
                    

# call main()
if __name__ == "__main__":
    main()


